# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Callback to print the information of MD simulation
"""

from ..optimizer import Updater


def print_info(step, potential, bias, optimizer, metrics):

    if_pbc = False
    use_updater = False
    _potential = potential.detach().cpu().numpy().squeeze()

    if isinstance(optimizer, Updater):
        temperature = optimizer.temperature
        kinetics = optimizer.kinetics
        use_updater = True
        _kinetics = kinetics.cpu().detach().numpy().sum(-1).squeeze()
        _temperature = temperature.cpu().detach().numpy().squeeze()
        _tot_energy = _potential + _kinetics
        if optimizer.system.pbc_box is not None:
            if_pbc = True
            pressure = optimizer.pressure
            _pressure = pressure.cpu().detach().numpy().sum(-1).squeeze()
            _volume = optimizer.system.get_volume().cpu().detach().numpy().squeeze()
        if bias is not None:
            _bias = bias.detach().cpu().numpy().squeeze()
        if metrics is not None:
            for k, v in metrics.items():
                metrics[k] = v.detach().cpu().numpy().squeeze()
    
    use_pbc = if_pbc

    if use_updater:
        info = 'Step: '+str(step) + ', '
        info += 'E_pot: ' + str(_potential) + ', '
        info += 'E_kin: ' + str(_kinetics) + ', '
        info += 'E_tot: ' + str(_tot_energy) + ', '
        info += 'Temperature: ' + str(_temperature)
        if use_pbc:
            info += ', '
            info += 'Pressure: ' + str(_pressure) + ', '
            info += 'Volume: ' + str(_volume)
        if bias is not None:
            info += ', '
            info += 'Bias: '+str(_bias)
        if metrics is not None:
            for k, v in metrics.items():
                info += f', {k}:{v}'
    else:
        info = 'Step: '+str(step) + ', '
        info += 'E_tot: ' + str(_potential) + ', '
    print('[TorchSPONGE]', info, flush=True)


